{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b7fc99c3",
   "metadata": {},
   "source": [
    "# Use QR-DQN to Play Pong\n",
    "\n",
    "TensorFlow version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2b23c9f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import copy\n",
    "import logging\n",
    "import itertools\n",
    "import sys\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)\n",
    "import pandas as pd\n",
    "import gym\n",
    "from gym.wrappers.atari_preprocessing import AtariPreprocessing\n",
    "from gym.wrappers.frame_stack import FrameStack\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow.compat.v2 as tf\n",
    "tf.random.set_seed(0)\n",
    "from tensorflow import keras\n",
    "from tensorflow import nn\n",
    "from tensorflow import optimizers\n",
    "from tensorflow import losses\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras import models\n",
    "\n",
    "logging.basicConfig(level=logging.DEBUG,\n",
    "        format='%(asctime)s [%(levelname)s] %(message)s',\n",
    "        stream=sys.stdout, datefmt='%H:%M:%S')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "693bdf1e",
   "metadata": {},
   "source": [
    "Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6ab1cb62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "05:01:52 [INFO] env: <AtariPreprocessing<TimeLimit<AtariEnv<PongNoFrameskip-v4>>>>\n",
      "05:01:52 [INFO] action_space: Discrete(6)\n",
      "05:01:52 [INFO] observation_space: Box(0, 255, (4, 84, 84), uint8)\n",
      "05:01:52 [INFO] reward_range: (-inf, inf)\n",
      "05:01:52 [INFO] metadata: {'render.modes': ['human', 'rgb_array']}\n",
      "05:01:52 [INFO] num_stack: 4\n",
      "05:01:52 [INFO] lz4_compress: False\n",
      "05:01:52 [INFO] frames: deque([], maxlen=4)\n",
      "05:01:52 [INFO] id: PongNoFrameskip-v4\n",
      "05:01:52 [INFO] entry_point: gym.envs.atari:AtariEnv\n",
      "05:01:52 [INFO] reward_threshold: None\n",
      "05:01:52 [INFO] nondeterministic: False\n",
      "05:01:52 [INFO] max_episode_steps: 400000\n",
      "05:01:52 [INFO] _kwargs: {'game': 'pong', 'obs_type': 'image', 'frameskip': 1}\n",
      "05:01:52 [INFO] _env_name: PongNoFrameskip\n"
     ]
    }
   ],
   "source": [
    "env = FrameStack(AtariPreprocessing(gym.make('PongNoFrameskip-v4')),\n",
    "        num_stack=4)\n",
    "for key in vars(env):\n",
    "    logging.info('%s: %s', key, vars(env)[key])\n",
    "for key in vars(env.spec):\n",
    "    logging.info('%s: %s', key, vars(env.spec)[key])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3747e919",
   "metadata": {},
   "source": [
    "Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5eb74be3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DQNReplayer:\n",
    "    def __init__(self, capacity):\n",
    "        self.memory = pd.DataFrame(index=range(capacity),\n",
    "                columns=['state', 'action', 'reward', 'next_state', 'done'])\n",
    "        self.i = 0\n",
    "        self.count = 0\n",
    "        self.capacity = capacity\n",
    "\n",
    "    def store(self, *args):\n",
    "        self.memory.loc[self.i] = args\n",
    "        self.i = (self.i + 1) % self.capacity\n",
    "        self.count = min(self.count + 1, self.capacity)\n",
    "\n",
    "    def sample(self, size):\n",
    "        indices = np.random.choice(self.count, size=size)\n",
    "        return (np.stack(self.memory.loc[indices, field]) for field in\n",
    "                self.memory.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a2b9a656",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Agent:\n",
    "    def __init__(self, env):\n",
    "        self.action_n = env.action_space.n\n",
    "        self.gamma = 0.99\n",
    "        self.epsilon = 1.\n",
    "\n",
    "        self.replayer = DQNReplayer(capacity=100000)\n",
    "\n",
    "        quantile_count = 200\n",
    "        self.cumprob_tensor = tf.range(1 / (2 * quantile_count),\n",
    "                1, 1 / quantile_count)[np.newaxis, :, np.newaxis]\n",
    "\n",
    "        self.evaluate_net = self.build_net(self.action_n, quantile_count)\n",
    "        self.target_net = models.clone_model(self.evaluate_net)\n",
    "\n",
    "    def build_net(self, action_n, quantile_count):\n",
    "        net = keras.Sequential([\n",
    "                keras.layers.Permute((2, 3, 1), input_shape=(4, 84, 84)),\n",
    "                layers.Conv2D(32, kernel_size=8, strides=4, activation=nn.relu),\n",
    "                layers.Conv2D(64, kernel_size=4, strides=2, activation=nn.relu),\n",
    "                layers.Conv2D(64, kernel_size=3, strides=1, activation=nn.relu),\n",
    "                layers.Flatten(),\n",
    "                layers.Dense(512, activation=nn.relu),\n",
    "                layers.Dense(action_n * quantile_count),\n",
    "                layers.Reshape((action_n, quantile_count))])\n",
    "        optimizer = optimizers.Adam(0.0001)\n",
    "        net.compile(optimizer=optimizer)\n",
    "        return net\n",
    "        \n",
    "    def reset(self, mode=None):\n",
    "        self.mode = mode\n",
    "        if mode == 'train':\n",
    "            self.trajectory = []\n",
    "\n",
    "    def step(self, observation, reward, done):\n",
    "        state_tensor = tf.convert_to_tensor(np.array(observation)[np.newaxis],\n",
    "                dtype=tf.float32)\n",
    "        q_component_tensor = self.evaluate_net(state_tensor)\n",
    "        q_tensor = tf.reduce_mean(q_component_tensor, axis=2)\n",
    "        action_tensor = tf.math.argmax(q_tensor, axis=1)\n",
    "        actions = action_tensor.numpy()\n",
    "        action = actions[0]\n",
    "        if self.mode == 'train':\n",
    "            if np.random.rand() < self.epsilon:\n",
    "                action = np.random.randint(0, self.action_n)\n",
    "            \n",
    "            self.trajectory += [observation, reward, done, action]\n",
    "            if len(self.trajectory) >= 8:\n",
    "                state, _, _, act, next_state, reward, done, _ = \\\n",
    "                        self.trajectory[-8:]\n",
    "                self.replayer.store(state, act, reward, next_state, done)\n",
    "            if self.replayer.count >= 1024 and self.replayer.count % 10 == 0:\n",
    "                self.learn()\n",
    "        return action\n",
    "\n",
    "    def close(self):\n",
    "        pass\n",
    "\n",
    "    def update_net(self, target_net, evaluate_net, learning_rate=0.005):\n",
    "        average_weights = [(1. - learning_rate) * t + learning_rate * e for t, e\n",
    "                in zip(target_net.get_weights(), evaluate_net.get_weights())]\n",
    "        target_net.set_weights(average_weights)\n",
    "\n",
    "    def learn(self):\n",
    "        # replay\n",
    "        batch_size = 32\n",
    "        states, actions, rewards, next_states, dones = \\\n",
    "                self.replayer.sample(batch_size)\n",
    "        state_tensor = tf.convert_to_tensor(states, dtype=tf.float32)\n",
    "        reward_tensor = tf.convert_to_tensor(rewards[:, np.newaxis],\n",
    "                dtype=tf.float32)\n",
    "        done_tensor = tf.convert_to_tensor(dones[:, np.newaxis],\n",
    "                dtype=tf.float32)\n",
    "        next_state_tensor = tf.convert_to_tensor(next_states, dtype=tf.float32)\n",
    "\n",
    "        # compute target\n",
    "        next_q_component_tensor = self.evaluate_net(next_state_tensor)\n",
    "        next_q_tensor = tf.reduce_mean(next_q_component_tensor, axis=2)\n",
    "        next_action_tensor = tf.math.argmax(next_q_tensor, axis=1)\n",
    "        next_actions = next_action_tensor.numpy()\n",
    "        all_next_q_quantile_tensor = self.target_net(next_state_tensor)\n",
    "        indices = [[idx, next_action] for idx, next_action in\n",
    "                enumerate(next_actions)]\n",
    "        next_q_quantile_tensor = tf.gather_nd(all_next_q_quantile_tensor,\n",
    "                indices)\n",
    "        target_quantile_tensor = reward_tensor + self.gamma \\\n",
    "                * next_q_quantile_tensor * (1. - done_tensor)\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            all_q_quantile_tensor = self.evaluate_net(state_tensor)\n",
    "            indices = [[idx, action] for idx, action in enumerate(actions)]\n",
    "            q_quantile_tensor = tf.gather_nd(all_q_quantile_tensor, indices)\n",
    "\n",
    "            target_quantile_tensor = target_quantile_tensor[:, np.newaxis, :]\n",
    "            q_quantile_tensor = q_quantile_tensor[:, :, np.newaxis]\n",
    "            td_error_tensor = target_quantile_tensor - q_quantile_tensor\n",
    "            abs_td_error_tensor = tf.math.abs(td_error_tensor)\n",
    "            hubor_delta = 1.\n",
    "            hubor_loss_tensor = tf.where(abs_td_error_tensor < hubor_delta,\n",
    "                    0.5 * tf.square(td_error_tensor),\n",
    "                    hubor_delta * (abs_td_error_tensor - 0.5 * hubor_delta))\n",
    "            comparison_tensor = tf.cast(td_error_tensor < 0, dtype=tf.float32)\n",
    "            quantile_regression_tensor = tf.math.abs(self.cumprob_tensor -\n",
    "                    comparison_tensor)\n",
    "            quantile_huber_loss_tensor = tf.reduce_mean(tf.reduce_sum(\n",
    "                    hubor_loss_tensor * quantile_regression_tensor, axis=-1),\n",
    "                    axis=1)\n",
    "            loss_tensor = tf.reduce_mean(quantile_huber_loss_tensor)\n",
    "        grads = tape.gradient(loss_tensor, self.evaluate_net.variables)\n",
    "        self.evaluate_net.optimizer.apply_gradients(\n",
    "                zip(grads, self.evaluate_net.variables))\n",
    "\n",
    "        self.update_net(self.target_net, self.evaluate_net)\n",
    "\n",
    "        self.epsilon = max(self.epsilon - 1e-5, 0.05)\n",
    "\n",
    "\n",
    "agent = Agent(env)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f02b76e5",
   "metadata": {},
   "source": [
    "Train & Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34ae8868",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "05:01:54 [INFO] ==== train ====\n",
      "05:02:09 [DEBUG] train episode 0: reward = -20.00, steps = 1049\n",
      "05:02:47 [DEBUG] train episode 1: reward = -19.00, steps = 1070\n",
      "05:03:18 [DEBUG] train episode 2: reward = -20.00, steps = 874\n",
      "05:03:55 [DEBUG] train episode 3: reward = -18.00, steps = 1031\n",
      "05:04:31 [DEBUG] train episode 4: reward = -20.00, steps = 1073\n",
      "05:04:57 [DEBUG] train episode 5: reward = -21.00, steps = 764\n",
      "05:05:34 [DEBUG] train episode 6: reward = -19.00, steps = 1114\n",
      "05:06:00 [DEBUG] train episode 7: reward = -21.00, steps = 778\n",
      "05:06:32 [DEBUG] train episode 8: reward = -19.00, steps = 958\n",
      "05:06:58 [DEBUG] train episode 9: reward = -21.00, steps = 762\n",
      "05:07:25 [DEBUG] train episode 10: reward = -21.00, steps = 821\n",
      "05:07:53 [DEBUG] train episode 11: reward = -21.00, steps = 823\n",
      "05:08:27 [DEBUG] train episode 12: reward = -19.00, steps = 1007\n",
      "05:09:01 [DEBUG] train episode 13: reward = -21.00, steps = 998\n",
      "05:09:33 [DEBUG] train episode 14: reward = -20.00, steps = 969\n",
      "05:10:00 [DEBUG] train episode 15: reward = -21.00, steps = 791\n",
      "05:10:26 [DEBUG] train episode 16: reward = -21.00, steps = 785\n",
      "05:11:02 [DEBUG] train episode 17: reward = -21.00, steps = 1063\n",
      "05:11:35 [DEBUG] train episode 18: reward = -20.00, steps = 959\n",
      "05:12:05 [DEBUG] train episode 19: reward = -21.00, steps = 898\n",
      "05:12:35 [DEBUG] train episode 20: reward = -21.00, steps = 906\n",
      "05:13:23 [DEBUG] train episode 21: reward = -15.00, steps = 1424\n",
      "05:13:55 [DEBUG] train episode 22: reward = -19.00, steps = 960\n",
      "05:14:21 [DEBUG] train episode 23: reward = -21.00, steps = 779\n",
      "05:14:56 [DEBUG] train episode 24: reward = -18.00, steps = 1023\n",
      "05:15:30 [DEBUG] train episode 25: reward = -20.00, steps = 1015\n",
      "05:15:59 [DEBUG] train episode 26: reward = -21.00, steps = 864\n",
      "05:16:28 [DEBUG] train episode 27: reward = -21.00, steps = 869\n",
      "05:17:08 [DEBUG] train episode 28: reward = -20.00, steps = 1197\n",
      "05:17:50 [DEBUG] train episode 29: reward = -18.00, steps = 1259\n",
      "05:18:25 [DEBUG] train episode 30: reward = -21.00, steps = 1023\n",
      "05:19:00 [DEBUG] train episode 31: reward = -20.00, steps = 1049\n",
      "05:19:26 [DEBUG] train episode 32: reward = -21.00, steps = 791\n",
      "05:19:57 [DEBUG] train episode 33: reward = -21.00, steps = 913\n",
      "05:20:32 [DEBUG] train episode 34: reward = -19.00, steps = 991\n",
      "05:21:04 [DEBUG] train episode 35: reward = -20.00, steps = 925\n",
      "05:21:34 [DEBUG] train episode 36: reward = -21.00, steps = 882\n",
      "05:22:09 [DEBUG] train episode 37: reward = -20.00, steps = 1038\n",
      "05:22:38 [DEBUG] train episode 38: reward = -20.00, steps = 868\n",
      "05:23:08 [DEBUG] train episode 39: reward = -20.00, steps = 898\n",
      "05:23:45 [DEBUG] train episode 40: reward = -19.00, steps = 1091\n",
      "05:24:25 [DEBUG] train episode 41: reward = -18.00, steps = 1154\n",
      "05:25:05 [DEBUG] train episode 42: reward = -19.00, steps = 1160\n",
      "05:25:36 [DEBUG] train episode 43: reward = -21.00, steps = 942\n",
      "05:26:08 [DEBUG] train episode 44: reward = -20.00, steps = 944\n",
      "05:26:36 [DEBUG] train episode 45: reward = -21.00, steps = 821\n",
      "05:27:07 [DEBUG] train episode 46: reward = -20.00, steps = 939\n",
      "05:27:34 [DEBUG] train episode 47: reward = -21.00, steps = 808\n",
      "05:28:03 [DEBUG] train episode 48: reward = -21.00, steps = 852\n",
      "05:28:38 [DEBUG] train episode 49: reward = -21.00, steps = 1010\n",
      "05:29:12 [DEBUG] train episode 50: reward = -19.00, steps = 1006\n",
      "05:29:46 [DEBUG] train episode 51: reward = -21.00, steps = 1000\n",
      "05:30:14 [DEBUG] train episode 52: reward = -21.00, steps = 808\n",
      "05:30:49 [DEBUG] train episode 53: reward = -21.00, steps = 1010\n",
      "05:31:18 [DEBUG] train episode 54: reward = -21.00, steps = 807\n",
      "05:31:49 [DEBUG] train episode 55: reward = -21.00, steps = 850\n",
      "05:32:20 [DEBUG] train episode 56: reward = -20.00, steps = 865\n",
      "05:32:57 [DEBUG] train episode 57: reward = -20.00, steps = 1019\n",
      "05:33:37 [DEBUG] train episode 58: reward = -18.00, steps = 1104\n",
      "05:34:10 [DEBUG] train episode 59: reward = -19.00, steps = 919\n",
      "05:34:45 [DEBUG] train episode 60: reward = -19.00, steps = 989\n",
      "05:35:19 [DEBUG] train episode 61: reward = -20.00, steps = 945\n",
      "05:35:51 [DEBUG] train episode 62: reward = -20.00, steps = 908\n",
      "05:36:27 [DEBUG] train episode 63: reward = -21.00, steps = 1000\n",
      "05:36:54 [DEBUG] train episode 64: reward = -21.00, steps = 759\n",
      "05:37:29 [DEBUG] train episode 65: reward = -20.00, steps = 973\n",
      "05:37:59 [DEBUG] train episode 66: reward = -21.00, steps = 844\n",
      "05:38:40 [DEBUG] train episode 67: reward = -21.00, steps = 1151\n",
      "05:39:12 [DEBUG] train episode 68: reward = -20.00, steps = 887\n",
      "05:39:48 [DEBUG] train episode 69: reward = -19.00, steps = 1011\n",
      "05:40:18 [DEBUG] train episode 70: reward = -20.00, steps = 840\n",
      "05:40:49 [DEBUG] train episode 71: reward = -20.00, steps = 864\n",
      "05:41:18 [DEBUG] train episode 72: reward = -21.00, steps = 782\n",
      "05:41:51 [DEBUG] train episode 73: reward = -21.00, steps = 941\n",
      "05:42:23 [DEBUG] train episode 74: reward = -20.00, steps = 885\n",
      "05:42:56 [DEBUG] train episode 75: reward = -20.00, steps = 903\n",
      "05:43:26 [DEBUG] train episode 76: reward = -20.00, steps = 835\n",
      "05:43:58 [DEBUG] train episode 77: reward = -20.00, steps = 885\n",
      "05:44:26 [DEBUG] train episode 78: reward = -21.00, steps = 786\n",
      "05:44:56 [DEBUG] train episode 79: reward = -21.00, steps = 847\n",
      "05:45:28 [DEBUG] train episode 80: reward = -20.00, steps = 899\n",
      "05:45:57 [DEBUG] train episode 81: reward = -21.00, steps = 820\n",
      "05:46:26 [DEBUG] train episode 82: reward = -21.00, steps = 820\n",
      "05:47:00 [DEBUG] train episode 83: reward = -20.00, steps = 974\n",
      "05:47:30 [DEBUG] train episode 84: reward = -20.00, steps = 859\n",
      "05:48:12 [DEBUG] train episode 85: reward = -18.00, steps = 1197\n",
      "05:48:41 [DEBUG] train episode 86: reward = -21.00, steps = 822\n",
      "05:49:14 [DEBUG] train episode 87: reward = -20.00, steps = 915\n",
      "05:49:51 [DEBUG] train episode 88: reward = -19.00, steps = 1024\n",
      "05:50:23 [DEBUG] train episode 89: reward = -20.00, steps = 882\n",
      "05:50:58 [DEBUG] train episode 90: reward = -21.00, steps = 961\n",
      "05:51:28 [DEBUG] train episode 91: reward = -21.00, steps = 823\n",
      "05:52:02 [DEBUG] train episode 92: reward = -19.00, steps = 936\n",
      "05:52:31 [DEBUG] train episode 93: reward = -21.00, steps = 794\n",
      "05:53:01 [DEBUG] train episode 94: reward = -21.00, steps = 821\n",
      "05:53:34 [DEBUG] train episode 95: reward = -21.00, steps = 895\n",
      "05:54:10 [DEBUG] train episode 96: reward = -21.00, steps = 972\n",
      "05:54:47 [DEBUG] train episode 97: reward = -19.00, steps = 999\n",
      "05:55:17 [DEBUG] train episode 98: reward = -21.00, steps = 819\n",
      "05:55:56 [DEBUG] train episode 99: reward = -19.00, steps = 1036\n",
      "05:56:27 [DEBUG] train episode 100: reward = -21.00, steps = 819\n",
      "05:57:05 [DEBUG] train episode 101: reward = -20.00, steps = 1009\n",
      "05:57:34 [DEBUG] train episode 102: reward = -21.00, steps = 778\n",
      "05:58:19 [DEBUG] train episode 103: reward = -20.00, steps = 1166\n",
      "05:58:59 [DEBUG] train episode 104: reward = -19.00, steps = 1085\n",
      "05:59:43 [DEBUG] train episode 105: reward = -18.00, steps = 1131\n",
      "06:02:38 [DEBUG] train episode 106: reward = -20.00, steps = 1025\n",
      "06:07:09 [DEBUG] train episode 107: reward = -19.00, steps = 1142\n",
      "06:10:58 [DEBUG] train episode 108: reward = -20.00, steps = 957\n",
      "06:15:50 [DEBUG] train episode 109: reward = -20.00, steps = 1198\n",
      "06:20:03 [DEBUG] train episode 110: reward = -20.00, steps = 1066\n",
      "06:24:16 [DEBUG] train episode 111: reward = -21.00, steps = 1060\n",
      "06:27:55 [DEBUG] train episode 112: reward = -19.00, steps = 934\n",
      "06:31:06 [DEBUG] train episode 113: reward = -21.00, steps = 811\n",
      "06:34:31 [DEBUG] train episode 114: reward = -20.00, steps = 882\n",
      "06:38:25 [DEBUG] train episode 115: reward = -21.00, steps = 1059\n",
      "06:42:29 [DEBUG] train episode 116: reward = -19.00, steps = 1099\n",
      "06:45:49 [DEBUG] train episode 117: reward = -21.00, steps = 897\n",
      "06:49:18 [DEBUG] train episode 118: reward = -21.00, steps = 938\n",
      "06:53:37 [DEBUG] train episode 119: reward = -18.00, steps = 1166\n",
      "06:57:52 [DEBUG] train episode 120: reward = -21.00, steps = 1153\n",
      "07:02:02 [DEBUG] train episode 121: reward = -19.00, steps = 1130\n",
      "07:06:32 [DEBUG] train episode 122: reward = -20.00, steps = 1226\n",
      "07:10:11 [DEBUG] train episode 123: reward = -21.00, steps = 985\n",
      "07:14:28 [DEBUG] train episode 124: reward = -20.00, steps = 1174\n",
      "07:18:16 [DEBUG] train episode 125: reward = -21.00, steps = 1046\n",
      "07:22:12 [DEBUG] train episode 126: reward = -19.00, steps = 1078\n",
      "07:26:11 [DEBUG] train episode 127: reward = -20.00, steps = 1096\n",
      "07:30:22 [DEBUG] train episode 128: reward = -19.00, steps = 1153\n",
      "07:35:31 [DEBUG] train episode 129: reward = -18.00, steps = 1403\n",
      "07:38:31 [DEBUG] train episode 130: reward = -21.00, steps = 819\n",
      "07:43:14 [DEBUG] train episode 131: reward = -19.00, steps = 1281\n",
      "07:47:35 [DEBUG] train episode 132: reward = -19.00, steps = 1183\n",
      "07:53:01 [DEBUG] train episode 133: reward = -16.00, steps = 1479\n",
      "07:56:40 [DEBUG] train episode 134: reward = -20.00, steps = 988\n",
      "08:00:59 [DEBUG] train episode 135: reward = -19.00, steps = 1173\n",
      "08:05:13 [DEBUG] train episode 136: reward = -20.00, steps = 1151\n",
      "08:10:13 [DEBUG] train episode 137: reward = -19.00, steps = 1358\n",
      "08:14:33 [DEBUG] train episode 138: reward = -19.00, steps = 1173\n",
      "08:18:40 [DEBUG] train episode 139: reward = -20.00, steps = 1124\n",
      "08:24:51 [DEBUG] train episode 140: reward = -19.00, steps = 1681\n",
      "08:29:51 [DEBUG] train episode 141: reward = -18.00, steps = 1359\n",
      "08:34:31 [DEBUG] train episode 142: reward = -19.00, steps = 1280\n",
      "08:39:34 [DEBUG] train episode 143: reward = -18.00, steps = 1376\n",
      "08:44:34 [DEBUG] train episode 144: reward = -18.00, steps = 1361\n",
      "08:50:48 [DEBUG] train episode 145: reward = -15.00, steps = 1699\n",
      "08:55:39 [DEBUG] train episode 146: reward = -17.00, steps = 1324\n",
      "09:03:40 [DEBUG] train episode 147: reward = -17.00, steps = 2199\n",
      "09:09:15 [DEBUG] train episode 148: reward = -19.00, steps = 1551\n",
      "09:17:21 [DEBUG] train episode 149: reward = -11.00, steps = 2208\n",
      "09:23:35 [DEBUG] train episode 150: reward = -17.00, steps = 1703\n",
      "09:29:54 [DEBUG] train episode 151: reward = -18.00, steps = 1715\n",
      "09:34:25 [DEBUG] train episode 152: reward = -20.00, steps = 1235\n",
      "09:40:46 [DEBUG] train episode 153: reward = -15.00, steps = 1732\n",
      "09:48:01 [DEBUG] train episode 154: reward = -13.00, steps = 1978\n",
      "09:53:48 [DEBUG] train episode 155: reward = -19.00, steps = 1598\n",
      "09:58:53 [DEBUG] train episode 156: reward = -19.00, steps = 1386\n",
      "10:05:47 [DEBUG] train episode 157: reward = -19.00, steps = 1889\n",
      "10:15:13 [DEBUG] train episode 158: reward = -17.00, steps = 2573\n",
      "10:21:28 [DEBUG] train episode 159: reward = -19.00, steps = 1696\n",
      "10:28:54 [DEBUG] train episode 160: reward = -18.00, steps = 2033\n",
      "10:36:04 [DEBUG] train episode 161: reward = -19.00, steps = 1948\n",
      "10:44:18 [DEBUG] train episode 162: reward = -17.00, steps = 2247\n",
      "10:51:43 [DEBUG] train episode 163: reward = -20.00, steps = 2045\n",
      "10:59:38 [DEBUG] train episode 164: reward = -18.00, steps = 2153\n",
      "11:08:05 [DEBUG] train episode 165: reward = -16.00, steps = 2201\n",
      "11:17:27 [DEBUG] train episode 166: reward = -16.00, steps = 2255\n",
      "11:27:47 [DEBUG] train episode 167: reward = -9.00, steps = 2516\n",
      "11:38:01 [DEBUG] train episode 168: reward = -10.00, steps = 2506\n",
      "11:42:37 [DEBUG] train episode 169: reward = -18.00, steps = 1049\n",
      "11:50:29 [DEBUG] train episode 170: reward = -11.00, steps = 1930\n",
      "11:58:11 [DEBUG] train episode 171: reward = -14.00, steps = 1857\n",
      "12:07:42 [DEBUG] train episode 172: reward = -7.00, steps = 2378\n",
      "12:17:09 [DEBUG] train episode 173: reward = -11.00, steps = 2399\n",
      "12:24:35 [DEBUG] train episode 174: reward = -15.00, steps = 1906\n",
      "12:32:47 [DEBUG] train episode 175: reward = -13.00, steps = 2031\n",
      "12:41:18 [DEBUG] train episode 176: reward = -13.00, steps = 2118\n",
      "12:50:20 [DEBUG] train episode 177: reward = -9.00, steps = 2332\n",
      "13:02:39 [DEBUG] train episode 178: reward = -5.00, steps = 3181\n",
      "13:12:22 [DEBUG] train episode 179: reward = -7.00, steps = 2470\n",
      "13:22:34 [DEBUG] train episode 180: reward = -6.00, steps = 2582\n",
      "13:33:27 [DEBUG] train episode 181: reward = -5.00, steps = 2816\n",
      "13:42:53 [DEBUG] train episode 182: reward = -5.00, steps = 2444\n",
      "13:53:59 [DEBUG] train episode 183: reward = -7.00, steps = 2887\n",
      "13:58:37 [DEBUG] train episode 184: reward = -19.00, steps = 1222\n",
      "14:05:59 [DEBUG] train episode 185: reward = -14.00, steps = 1916\n",
      "14:18:17 [DEBUG] train episode 186: reward = 1.00, steps = 3129\n",
      "14:28:27 [DEBUG] train episode 187: reward = -7.00, steps = 2408\n",
      "14:38:43 [DEBUG] train episode 188: reward = -3.00, steps = 2506\n",
      "14:46:11 [DEBUG] train episode 189: reward = -14.00, steps = 1818\n",
      "14:57:59 [DEBUG] train episode 190: reward = -4.00, steps = 2806\n",
      "15:05:57 [DEBUG] train episode 191: reward = -12.00, steps = 1921\n",
      "15:12:53 [DEBUG] train episode 192: reward = -13.00, steps = 1728\n",
      "15:23:13 [DEBUG] train episode 193: reward = -3.00, steps = 2657\n",
      "15:33:05 [DEBUG] train episode 194: reward = -5.00, steps = 2523\n",
      "15:41:17 [DEBUG] train episode 195: reward = -12.00, steps = 2119\n",
      "15:51:03 [DEBUG] train episode 196: reward = -6.00, steps = 2476\n",
      "16:01:02 [DEBUG] train episode 197: reward = -3.00, steps = 2544\n",
      "16:09:19 [DEBUG] train episode 198: reward = -10.00, steps = 2063\n",
      "16:19:18 [DEBUG] train episode 199: reward = -4.00, steps = 2539\n",
      "16:28:32 [DEBUG] train episode 200: reward = -5.00, steps = 2374\n",
      "16:36:41 [DEBUG] train episode 201: reward = -9.00, steps = 1945\n",
      "16:45:31 [DEBUG] train episode 202: reward = -7.00, steps = 2124\n",
      "16:54:43 [DEBUG] train episode 203: reward = -8.00, steps = 2221\n",
      "17:00:09 [DEBUG] train episode 204: reward = -16.00, steps = 1309\n",
      "17:09:55 [DEBUG] train episode 205: reward = -9.00, steps = 2360\n",
      "17:15:15 [DEBUG] train episode 206: reward = -16.00, steps = 1284\n",
      "17:23:18 [DEBUG] train episode 207: reward = -11.00, steps = 1966\n",
      "17:32:59 [DEBUG] train episode 208: reward = -6.00, steps = 2368\n",
      "17:44:07 [DEBUG] train episode 209: reward = -3.00, steps = 2688\n",
      "17:52:58 [DEBUG] train episode 210: reward = -8.00, steps = 2047\n",
      "18:01:49 [DEBUG] train episode 211: reward = -7.00, steps = 2126\n",
      "18:10:11 [DEBUG] train episode 212: reward = -10.00, steps = 2061\n",
      "18:19:15 [DEBUG] train episode 213: reward = -6.00, steps = 2308\n",
      "18:29:10 [DEBUG] train episode 214: reward = -3.00, steps = 2558\n",
      "18:37:44 [DEBUG] train episode 215: reward = -7.00, steps = 2133\n",
      "18:46:41 [DEBUG] train episode 216: reward = -7.00, steps = 2156\n",
      "18:55:43 [DEBUG] train episode 217: reward = -6.00, steps = 2190\n",
      "19:04:11 [DEBUG] train episode 218: reward = -8.00, steps = 2128\n",
      "19:12:56 [DEBUG] train episode 219: reward = -6.00, steps = 2236\n",
      "19:21:53 [DEBUG] train episode 220: reward = -7.00, steps = 2195\n",
      "19:32:37 [DEBUG] train episode 221: reward = -4.00, steps = 2568\n",
      "19:42:52 [DEBUG] train episode 222: reward = -4.00, steps = 2500\n",
      "19:49:08 [DEBUG] train episode 223: reward = -14.00, steps = 1524\n",
      "19:57:24 [DEBUG] train episode 224: reward = -8.00, steps = 2084\n",
      "20:04:36 [DEBUG] train episode 225: reward = -11.00, steps = 1776\n",
      "20:13:43 [DEBUG] train episode 226: reward = -7.00, steps = 2256\n",
      "20:22:54 [DEBUG] train episode 227: reward = -8.00, steps = 2273\n",
      "20:30:08 [DEBUG] train episode 228: reward = -11.00, steps = 1836\n",
      "20:38:01 [DEBUG] train episode 229: reward = -10.00, steps = 1906\n",
      "20:47:15 [DEBUG] train episode 230: reward = -8.00, steps = 2263\n",
      "20:56:36 [DEBUG] train episode 231: reward = -6.00, steps = 2381\n",
      "21:07:31 [DEBUG] train episode 232: reward = -4.00, steps = 2759\n",
      "21:16:29 [DEBUG] train episode 233: reward = -8.00, steps = 2227\n",
      "21:23:40 [DEBUG] train episode 234: reward = -11.00, steps = 1805\n",
      "21:33:28 [DEBUG] train episode 235: reward = -4.00, steps = 2450\n",
      "21:43:03 [DEBUG] train episode 236: reward = -8.00, steps = 2355\n",
      "21:53:07 [DEBUG] train episode 237: reward = -6.00, steps = 2407\n",
      "22:04:26 [DEBUG] train episode 238: reward = -3.00, steps = 2682\n",
      "22:13:43 [DEBUG] train episode 239: reward = -4.00, steps = 2347\n",
      "22:24:28 [DEBUG] train episode 240: reward = -3.00, steps = 2670\n",
      "22:34:58 [DEBUG] train episode 241: reward = -1.00, steps = 2542\n",
      "22:41:04 [DEBUG] train episode 242: reward = -13.00, steps = 1531\n",
      "22:51:24 [DEBUG] train episode 243: reward = -5.00, steps = 2504\n",
      "23:01:23 [DEBUG] train episode 244: reward = -3.00, steps = 2523\n",
      "23:09:44 [DEBUG] train episode 245: reward = -8.00, steps = 2050\n",
      "23:20:00 [DEBUG] train episode 246: reward = -4.00, steps = 2577\n",
      "23:28:01 [DEBUG] train episode 247: reward = -8.00, steps = 1960\n",
      "23:36:19 [DEBUG] train episode 248: reward = -8.00, steps = 2045\n",
      "23:46:17 [DEBUG] train episode 249: reward = -2.00, steps = 2529\n",
      "23:56:39 [DEBUG] train episode 250: reward = -7.00, steps = 2618\n",
      "00:06:32 [DEBUG] train episode 251: reward = 5.00, steps = 2508\n",
      "00:15:29 [DEBUG] train episode 252: reward = -4.00, steps = 2297\n",
      "00:25:48 [DEBUG] train episode 253: reward = -1.00, steps = 2510\n",
      "00:35:36 [DEBUG] train episode 254: reward = -5.00, steps = 2347\n",
      "00:44:35 [DEBUG] train episode 255: reward = -5.00, steps = 2136\n",
      "00:55:00 [DEBUG] train episode 256: reward = -3.00, steps = 2594\n",
      "01:02:23 [DEBUG] train episode 257: reward = 13.00, steps = 1980\n",
      "01:11:38 [DEBUG] train episode 258: reward = -5.00, steps = 2476\n",
      "01:20:50 [DEBUG] train episode 259: reward = -1.00, steps = 2458\n",
      "01:30:07 [DEBUG] train episode 260: reward = 7.00, steps = 2494\n",
      "01:37:47 [DEBUG] train episode 261: reward = -6.00, steps = 2069\n",
      "01:46:36 [DEBUG] train episode 262: reward = -5.00, steps = 2373\n",
      "01:56:10 [DEBUG] train episode 263: reward = 4.00, steps = 2567\n",
      "02:03:01 [DEBUG] train episode 264: reward = 17.00, steps = 1822\n",
      "02:11:38 [DEBUG] train episode 265: reward = 11.00, steps = 2306\n",
      "02:18:43 [DEBUG] train episode 266: reward = 17.00, steps = 1908\n",
      "02:27:30 [DEBUG] train episode 267: reward = 5.00, steps = 2352\n",
      "02:34:50 [DEBUG] train episode 268: reward = -8.00, steps = 1972\n",
      "02:44:24 [DEBUG] train episode 269: reward = 3.00, steps = 2580\n",
      "02:51:15 [DEBUG] train episode 270: reward = 17.00, steps = 1846\n",
      "02:58:24 [DEBUG] train episode 271: reward = -10.00, steps = 1928\n",
      "03:06:36 [DEBUG] train episode 272: reward = 8.00, steps = 2204\n",
      "03:14:33 [DEBUG] train episode 273: reward = 12.00, steps = 2146\n",
      "03:22:18 [DEBUG] train episode 274: reward = 12.00, steps = 2085\n",
      "03:31:15 [DEBUG] train episode 275: reward = 6.00, steps = 2368\n",
      "03:38:59 [DEBUG] train episode 276: reward = -6.00, steps = 2087\n",
      "03:47:39 [DEBUG] train episode 277: reward = -5.00, steps = 2324\n",
      "03:57:14 [DEBUG] train episode 278: reward = 1.00, steps = 2572\n",
      "04:06:56 [DEBUG] train episode 279: reward = 1.00, steps = 2595\n",
      "04:14:55 [DEBUG] train episode 280: reward = 13.00, steps = 2143\n",
      "04:24:15 [DEBUG] train episode 281: reward = -1.00, steps = 2503\n",
      "04:32:46 [DEBUG] train episode 282: reward = 8.00, steps = 2284\n",
      "04:41:43 [DEBUG] train episode 283: reward = 1.00, steps = 2405\n",
      "04:49:59 [DEBUG] train episode 284: reward = 13.00, steps = 2216\n",
      "04:58:20 [DEBUG] train episode 285: reward = 8.00, steps = 2240\n",
      "05:05:53 [DEBUG] train episode 286: reward = 12.00, steps = 2027\n",
      "05:12:48 [DEBUG] train episode 287: reward = 16.00, steps = 1861\n",
      "05:20:12 [DEBUG] train episode 288: reward = 13.00, steps = 1995\n",
      "05:30:05 [DEBUG] train episode 289: reward = -2.00, steps = 2653\n",
      "05:39:11 [DEBUG] train episode 290: reward = 8.00, steps = 2377\n",
      "05:48:30 [DEBUG] train episode 291: reward = 8.00, steps = 2352\n",
      "05:59:05 [DEBUG] train episode 292: reward = 1.00, steps = 2751\n",
      "06:09:02 [DEBUG] train episode 293: reward = -1.00, steps = 2657\n",
      "06:19:18 [DEBUG] train episode 294: reward = 1.00, steps = 2741\n",
      "06:26:58 [DEBUG] train episode 295: reward = 13.00, steps = 2045\n",
      "06:37:03 [DEBUG] train episode 296: reward = -1.00, steps = 2693\n",
      "06:46:46 [DEBUG] train episode 297: reward = -1.00, steps = 2595\n",
      "06:54:27 [DEBUG] train episode 298: reward = 12.00, steps = 2048\n",
      "07:05:24 [DEBUG] train episode 299: reward = 1.00, steps = 2867\n",
      "07:14:47 [DEBUG] train episode 300: reward = 8.00, steps = 2503\n",
      "07:24:55 [DEBUG] train episode 301: reward = -1.00, steps = 2707\n",
      "07:33:04 [DEBUG] train episode 302: reward = -6.00, steps = 2176\n",
      "07:43:07 [DEBUG] train episode 303: reward = -1.00, steps = 2691\n",
      "07:53:02 [DEBUG] train episode 304: reward = -2.00, steps = 2661\n",
      "08:02:57 [DEBUG] train episode 305: reward = -2.00, steps = 2670\n",
      "08:10:29 [DEBUG] train episode 306: reward = 12.00, steps = 2031\n",
      "08:17:19 [DEBUG] train episode 307: reward = 16.00, steps = 1838\n",
      "08:26:10 [DEBUG] train episode 308: reward = 2.00, steps = 2369\n",
      "08:33:52 [DEBUG] train episode 309: reward = -6.00, steps = 2070\n",
      "08:43:22 [DEBUG] train episode 310: reward = -1.00, steps = 2556\n",
      "08:52:00 [DEBUG] train episode 311: reward = 7.00, steps = 2327\n",
      "08:58:53 [DEBUG] train episode 312: reward = -10.00, steps = 1848\n",
      "09:07:42 [DEBUG] train episode 313: reward = 7.00, steps = 2374\n",
      "09:14:52 [DEBUG] train episode 314: reward = 17.00, steps = 1925\n",
      "09:24:29 [DEBUG] train episode 315: reward = 5.00, steps = 2575\n",
      "09:32:23 [DEBUG] train episode 316: reward = -9.00, steps = 2126\n",
      "09:40:54 [DEBUG] train episode 317: reward = 8.00, steps = 2295\n",
      "09:49:25 [DEBUG] train episode 318: reward = 8.00, steps = 2280\n",
      "09:58:06 [DEBUG] train episode 319: reward = 13.00, steps = 2256\n",
      "10:07:00 [DEBUG] train episode 320: reward = 10.00, steps = 2229\n",
      "10:16:07 [DEBUG] train episode 321: reward = 13.00, steps = 2241\n",
      "10:23:41 [DEBUG] train episode 322: reward = -11.00, steps = 1856\n",
      "10:32:34 [DEBUG] train episode 323: reward = 6.00, steps = 2221\n",
      "10:40:36 [DEBUG] train episode 324: reward = -7.00, steps = 1974\n",
      "10:50:23 [DEBUG] train episode 325: reward = 7.00, steps = 2400\n",
      "10:57:48 [DEBUG] train episode 326: reward = 18.00, steps = 1778\n",
      "11:03:43 [DEBUG] train episode 327: reward = -14.00, steps = 1485\n",
      "11:12:20 [DEBUG] train episode 328: reward = 13.00, steps = 2107\n",
      "11:20:42 [DEBUG] train episode 329: reward = 13.00, steps = 2044\n",
      "11:29:37 [DEBUG] train episode 330: reward = 8.00, steps = 2178\n",
      "11:38:35 [DEBUG] train episode 331: reward = 8.00, steps = 2270\n",
      "11:48:34 [DEBUG] train episode 332: reward = 1.00, steps = 2566\n",
      "11:56:52 [DEBUG] train episode 333: reward = 10.00, steps = 2129\n",
      "12:05:16 [DEBUG] train episode 334: reward = 13.00, steps = 2157\n",
      "12:16:19 [DEBUG] train episode 335: reward = -1.00, steps = 2824\n"
     ]
    }
   ],
   "source": [
    "def play_episode(env, agent, max_episode_steps=None, mode=None, render=False):\n",
    "    observation, reward, done = env.reset(), 0., False\n",
    "    agent.reset(mode=mode)\n",
    "    episode_reward, elapsed_steps = 0., 0\n",
    "    while True:\n",
    "        action = agent.step(observation, reward, done)\n",
    "        if render:\n",
    "            env.render()\n",
    "        if done:\n",
    "            break\n",
    "        observation, reward, done, _ = env.step(action)\n",
    "        episode_reward += reward\n",
    "        elapsed_steps += 1\n",
    "        if max_episode_steps and elapsed_steps >= max_episode_steps:\n",
    "            break\n",
    "    agent.close()\n",
    "    return episode_reward, elapsed_steps\n",
    "\n",
    "\n",
    "logging.info('==== train ====')\n",
    "episode_rewards = []\n",
    "for episode in itertools.count():\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent, mode='train')\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('train episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "    if np.mean(episode_rewards[-5:]) > 16.:\n",
    "        break\n",
    "plt.plot(episode_rewards)\n",
    "\n",
    "\n",
    "logging.info('==== test ====')\n",
    "episode_rewards = []\n",
    "for episode in range(100):\n",
    "    episode_reward, elapsed_steps = play_episode(env, agent)\n",
    "    episode_rewards.append(episode_reward)\n",
    "    logging.debug('test episode %d: reward = %.2f, steps = %d',\n",
    "            episode, episode_reward, elapsed_steps)\n",
    "logging.info('average episode reward = %.2f ± %.2f',\n",
    "        np.mean(episode_rewards), np.std(episode_rewards))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
